import pandas as pd
import numpy as np
from torch.utils.data import TensorDataset, DataLoader
import torch


def prepare_data_loaders():
    train_dataset_db = pd.read_csv("../datapreprocessresult/reduced_data.csv", index_col=None)
    train_dataset_gotmd = pd.read_csv("../datapreprocessresult/reduced_gotmd.csv", index_col=None)
    train_dataset = pd.concat([train_dataset_db, train_dataset_gotmd], ignore_index=True)
    train_dataset = train_dataset.iloc[:276352]
    # train_dataset = train_dataset.iloc[:512]
    train_dataset = train_dataset.sample(n=8192, random_state=42)
    features = train_dataset.drop(['label', 'dborgot'], axis=1)
    labels = train_dataset['label']
    dborgotlabels = train_dataset['dborgot']
    features = np.array(features)
    labels = np.array(labels)
    dborgotlabels = np.array(dborgotlabels)
    features = torch.tensor(features, dtype=torch.float32)
    labels = torch.tensor(labels, dtype=torch.float32)
    dborgotlabels = torch.tensor(dborgotlabels, dtype=torch.float32)

    dataset = TensorDataset(features, labels, dborgotlabels)
    # print(dataset.__len__())
    # exit()
    train_loader = DataLoader(dataset=dataset,
                              batch_size=128,
                              shuffle=True)
    test_loader = DataLoader(dataset=dataset,
                             batch_size=128,
                             shuffle=True)

    # print(train_dataset_db.head())
    return train_loader, test_loader
